#include "defectrecognizer.h"
#include <QDebug>
#include "../mainwindow.h"
#include "../parameter/cnninfo.h"
#include "../parameter/imageinfo.h"

DefectRecognizer::DefectRecognizer() :
    m_bUseGPU(false),
    m_nImageHeight(3672),
    m_nImageWidth(5496),
    m_nPartHeight(640),
    m_nPartWidth(640)
{
    AdjustPartSize();
    InitYoloParameter();
}

void DefectRecognizer::InitYoloParameter(){
    // 边框颜色
    srand(time(0));
    for (int i = 0; i < 80; i++) {
        int b = rand() % 256;
        int g = rand() % 256;
        int r = rand() % 256;
        m_ObjColor.push_back(cv::Scalar(b, g, r));
    }
    LoadYoloModel();
}

void DefectRecognizer::LoadYoloModel(){
    CnnInfo& cnnInfo = MainWindow::instance()->GetSeatDetection()->GetCnnInfo();
    ImageInfo& imageInfo = MainWindow::instance()->GetSeatDetection()->GetImageInfo();
    m_nImageWidth = imageInfo.GetImageWidth();
    m_nImageHeight = imageInfo.GetImageHeight();
    m_nPartWidth = cnnInfo.GetPartWidth();
    m_nPartHeight = cnnInfo.GetPartHeight();

    //std::string strModelPath = cnnInfo.GetONNXModelPath().toStdString();
    //std::string str = "D:/test/seat/yolov8n.onnx";
    std::string strModelPath = m_strYoloOnnxModelPath.toStdString();
    bool isOK = m_objYolov8.readModel(m_objNet3, strModelPath, m_bUseGPU);

    if (isOK) {
        qDebug() << "Read YOLO onnx model ok!";
    }
    else {
        qDebug() << "Read YOLO onnx model failed!";
        return ;
    }
}

void DefectRecognizer::SetImageHeight(int value){
    m_nImageHeight = value;
    AdjustPartSize();
}

void DefectRecognizer::SetImageWidth(int value){
    m_nImageWidth = value;
    AdjustPartSize();
}

void DefectRecognizer::SetPartHeight(int value){
    m_nPartHeight = value;
    AdjustPartSize();
}

void DefectRecognizer::SetPartWidth(int value){
    m_nPartWidth = value;
    AdjustPartSize();
}

void DefectRecognizer::AdjustPartSize(){
    m_nCountHeight = m_nImageHeight / m_nPartHeight + (m_nImageHeight % m_nPartHeight == 0 ? 0 :1);
    m_nCountWidth = m_nImageWidth / m_nPartWidth + (m_nImageWidth % m_nPartWidth == 0 ? 0 :1);
}

void DefectRecognizer::SetUseGPU(bool useGPU){
    m_bUseGPU = useGPU;
}

std::vector<Detection> DefectRecognizer::RecognizePartsDefect(cv::Mat image){
    std::vector<Detection> results;
    for(int row = 0 ; row < m_nCountHeight; row ++){
        for(int col = 0; col < m_nCountWidth; col++){
            int w = ((col+1) * m_nPartWidth > m_nImageWidth) ? (m_nImageWidth % m_nPartWidth) : m_nPartWidth;
            int h = ((row+1) * m_nPartHeight > m_nImageHeight) ? (m_nImageHeight % m_nPartHeight) : m_nPartHeight;

            cv::Rect roi(col * m_nPartWidth,row * m_nPartHeight,w,h);
            cv::Mat subImage = image(roi);
            std::vector<Detection> result = m_objYolov8.Detect(subImage, m_objNet3);
            // 将 vector2 的元素添加到 vector1 的末尾
            for (int i = 0; i < result.size(); ++i) {
                // 调整块坐标
                result[i].box.x += col * m_nPartWidth;
                result[i].box.y += row * m_nPartHeight;

                results.push_back(result[i]);
            }
        }
    }

    return results;
}

void DefectRecognizer::DrawPred(cv::Mat image,std::vector<Detection> results){
    m_objYolov8.drawPred(image, results, m_ObjColor);
}

void DefectRecognizer::test(){
    std::vector<Detection> results;
    cv::Mat imageSrc = cv::imread("d:/test/defect/seatface0.bmp");
    cv::Mat imageDst = imageSrc.clone();

    qint64 CurrentMEpoch1 = QDateTime::currentMSecsSinceEpoch();
    //std::vector<Detection> result = m_objYolov8.Detect(imageDst, m_objNet3);
    for(int row = 0 ; row < m_nCountHeight; row ++){
        for(int col = 0; col < m_nCountWidth; col++){
            int w = ((col+1) * m_nPartWidth > m_nImageWidth) ? (m_nImageWidth % m_nPartWidth) : m_nPartWidth;
            int h = ((row+1) * m_nPartHeight > m_nImageHeight) ? (m_nImageHeight % m_nPartHeight) : m_nPartHeight;

            cv::Rect roi(col * m_nPartWidth,row * m_nPartHeight,w,h);
            cv::Mat subImage = imageSrc(roi);
            std::vector<Detection> result = m_objYolov8.Detect(subImage, m_objNet3);
            // 将 vector2 的元素添加到 vector1 的末尾
            for (int i = 0; i < result.size(); ++i) {
                // 调整块坐标
                result[i].box.x += col * m_nPartWidth;
                result[i].box.y += row * m_nPartHeight;

                results.push_back(result[i]);
            }
        }
    }

    qint64 CurrentMEpoch2 = QDateTime::currentMSecsSinceEpoch();
    QString elapse = QString::number(CurrentMEpoch2 - CurrentMEpoch1);

    qDebug() << elapse << "ms";

    m_objYolov8.drawPred(imageDst, results, m_ObjColor);

    // 保存到指定的文件夹下
    imwrite("d:/test/defect/box_yolov8.jpg", imageDst);
}
